package com.dmjy.pub.export.data2csv.service;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.text.csv.CsvUtil;
import cn.hutool.core.text.csv.CsvWriter;
import cn.hutool.core.util.StrUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.io.File;
import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

/**
 * @author 敖癸
 * @formatter:on
 * @since 2024/3/7
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ExportDataService {

    private final JdbcTemplate jdbcTemplate;

    @Async
    public void exportDataToFile(String sql, File file) {
        long beginTime = System.currentTimeMillis();
        AtomicInteger readCnt = new AtomicInteger();
        try (CsvWriter csvWriter = CsvUtil.getWriter(file, Charset.defaultCharset())) {
            jdbcTemplate.query(conn -> buildStmt(conn, sql), rs -> {
                handlerRow(rs, csvWriter, readCnt);
            });
            csvWriter.flush();
        }
        log.info("export complete, total: {}, costTime: {} MS", readCnt.get(), System.currentTimeMillis() - beginTime);
        // 全部导出完成后, 重命名文件, 把.ing后缀去掉
        FileUtil.rename(file, StrUtil.subBefore(file.getName(), ".ing", true), true);
    }

    /**
     * 构建预执行SQL模板
     *
     * @param conn 数据连接
     * @param sql  SQL
     * @return java.sql.PreparedStatement
     * @author 敖癸
     * @since 2024/3/15 - 15:37
     */
    private static PreparedStatement buildStmt(Connection conn, String sql) throws SQLException {
        PreparedStatement stmt = conn.prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
        // 固定写法, 因为mysql比较特殊, 要开启游标查询, 需要设置fetchSize为Integer.MIN_VALUE。其它数据库根据实际需求来调整大小
        stmt.setFetchSize(Integer.MIN_VALUE);
        stmt.setFetchDirection(ResultSet.FETCH_FORWARD);
        return stmt;
    }

    /** 处理每一行的数据 */
    private static void handlerRow(ResultSet rs, CsvWriter csvWriter, AtomicInteger readCnt) throws SQLException {
        // 获取每一行的字段数量
        int columnCount = rs.getMetaData().getColumnCount();
        // 按顺序获取每个字段的值
        String[] fields = IntStream.rangeClosed(1, columnCount).mapToObj(i -> getFieldValue(rs, i)).toArray(String[]::new);
        // 将一行数据append写入到csv文件
        csvWriter.writeLine(fields);
        // 每5000行清空一次缓冲区, 将数据持久化到磁盘
        if (readCnt.incrementAndGet() % 5000 == 0) {
            csvWriter.flush();
            log.info("exported: {}", readCnt.get());
        }
    }

    private static String getFieldValue(ResultSet rs, int i) {
        try {
            return rs.getString(i);
        } catch (SQLException e) {
            log.warn("field parsing failure", e);
            return "N/A";
        }
    }

}
